srun --job-name=raft -p xxx -n1 --gres=gpu:2 --mpi=pmi2 python -m torch.distributed.launch  --nproc_per_node=2 --master_port 20111  train.py --name raft-chairs --stage chairs --validation chairs --num_steps 100000 --batch_size 4 --lr 0.0004 --image_size 368 496 --wdecay 0.0001
